#include "DirectX11.h"

HRESULT NewDirectX11Session(DirectX11Session *session, UINT display, bool rgb)
{
    HRESULT hr = S_OK;
    UINT count = 0;

    D3D_FEATURE_LEVEL featureLevel;
    DXGI_OUTDUPL_FRAME_INFO frameInfo;
    DXGI_OUTPUT_DESC desc;
    DXGI_OUTDUPL_DESC duplDesc;

    IDXGIFactory1 *factory = NULL;
    IDXGIAdapter *adapter = NULL;
    IDXGIOutput *output = NULL;
    IDXGIOutput1 *output1 = NULL;
    IDXGIResource *resource = NULL;

    session->rgb = rgb;

    hr = CreateDXGIFactory1(&IID_IDXGIFactory, (void **)&factory);
    if (FAILED(hr))
    {
        goto release;
    }

    for (UINT i = 0;; i++)
    {
        hr = factory->lpVtbl->EnumAdapters(factory, i, &adapter);
        if (hr == DXGI_ERROR_NOT_FOUND)
        {
            break;
        }
        if (FAILED(hr))
        {
            goto release;
        }

        for (UINT j = 0;; j++)
        {
            hr = adapter->lpVtbl->EnumOutputs(adapter, j, &output);
            if (hr == DXGI_ERROR_NOT_FOUND)
            {
                break;
            }
            if (FAILED(hr))
            {
                goto release;
            }

            output->lpVtbl->GetDesc(output, &desc);
            if (desc.AttachedToDesktop)
            {
                if (count == display)
                {
                    break;
                }
                else
                {
                    count++;
                }
            }
            else
            {
                output->lpVtbl->Release(output);
                output = NULL;
            }
        }

        if (output != NULL)
        {
            break;
        }
        else
        {
            adapter->lpVtbl->Release(adapter);
            adapter = NULL;
        }
    }

    if (adapter == NULL || output == NULL)
    {
        hr = DXGI_ERROR_NOT_FOUND;
        goto release;
    }

    hr = D3D11CreateDevice(adapter, D3D_DRIVER_TYPE_UNKNOWN, NULL, D3D11_CREATE_DEVICE_BGRA_SUPPORT | D3D11_CREATE_DEVICE_VIDEO_SUPPORT, FeatureLevels, NumFeatureLevels, D3D11_SDK_VERSION, &session->device, &featureLevel, &session->context);
    if (FAILED(hr))
    {
        goto release;
    }

    hr = CompileShader(session);
    if (FAILED(hr))
    {
        goto release;
    }

    hr = output->lpVtbl->QueryInterface(output, &IID_IDXGIOutput1, (void **)&output1);
    if (FAILED(hr))
    {
        goto release;
    }

    output1->lpVtbl->GetDesc(output1, &duplDesc);

    hr = output1->lpVtbl->DuplicateOutput(output1, (IUnknown *)session->device, &session->duplication);
    if (FAILED(hr))
    {
        goto release;
    }
    if (session->duplication == NULL || session->duplication->lpVtbl == NULL)
    {
        hr = E_UNEXPECTED;
        goto release;
    }

release:
    if (factory != NULL && factory->lpVtbl != NULL)
    {
        factory->lpVtbl->Release(factory);
        factory = NULL;
    }
    if (adapter != NULL && adapter->lpVtbl != NULL)
    {
        adapter->lpVtbl->Release(adapter);
        adapter = NULL;
    }
    if (output != NULL && output->lpVtbl != NULL)
    {
        output->lpVtbl->Release(output);
        output = NULL;
    }
    if (output1 != NULL && output1->lpVtbl != NULL)
    {
        output1->lpVtbl->Release(output1);
        output1 = NULL;
    }
    if (resource != NULL && resource->lpVtbl != NULL)
    {
        resource->lpVtbl->Release(resource);
        resource = NULL;
    }

    return hr;
}

void CloseDirectX11Session(DirectX11Session *session)
{
    if (session->duplication != NULL && session->duplication->lpVtbl != NULL)
    {
        session->duplication->lpVtbl->Release(session->duplication);
        session->duplication = NULL;
    }
    if (session->context != NULL && session->context->lpVtbl != NULL)
    {
        session->context->lpVtbl->Release(session->context);
        session->context = NULL;
    }
    if (session->device != NULL && session->device->lpVtbl != NULL)
    {
        session->device->lpVtbl->Release(session->device);
        session->device = NULL;
    }
    if (session->shader != NULL && session->shader->lpVtbl != NULL)
    {
        session->shader->lpVtbl->Release(session->shader);
        session->shader = NULL;
    }
}

HRESULT Screenshot(DirectX11Session *session, INT x, INT y, UINT *width, UINT *height, BYTE **bits)
{
    HRESULT hr = S_OK;

    DXGI_OUTDUPL_FRAME_INFO frameInfo;
    D3D11_TEXTURE2D_DESC desc;
    D3D11_MAPPED_SUBRESOURCE mapped;

    IDXGIResource *resource = NULL;
    ID3D11Texture2D *texture = NULL;
    ID3D11Texture2D *shade = NULL;
    ID3D11Texture2D *staging = NULL;
    ID3D11ShaderResourceView *view = NULL;
    ID3D11UnorderedAccessView *uav = NULL;

    hr = session->duplication->lpVtbl->AcquireNextFrame(session->duplication, 0, &frameInfo, &resource); // INFINITE
    if (FAILED(hr))
    {
        goto release_frame;
    }

    hr = resource->lpVtbl->QueryInterface(resource, &IID_ID3D11Texture2D, (void **)&texture);
    if (FAILED(hr))
    {
        goto release_frame;
    }

    texture->lpVtbl->GetDesc(texture, &desc);

    if (session->rgb != NULL && session->rgb)
    {
        desc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
        desc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_UNORDERED_ACCESS;
        desc.MipLevels = 1;
        desc.MiscFlags = 0;
        hr = session->device->lpVtbl->CreateTexture2D(session->device, &desc, NULL, &shade);
        if (FAILED(hr))
        {
            goto release_frame;
        }
        D3D11_SHADER_RESOURCE_VIEW_DESC viewDesc = {
            .Format = desc.Format,
            .ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D,
            .Texture2D = {.MostDetailedMip = 0, .MipLevels = desc.MipLevels}};
        hr = session->device->lpVtbl->CreateShaderResourceView(session->device, (ID3D11Resource *)texture, &viewDesc, &view);
        if (FAILED(hr))
        {
            goto release_frame;
        }
        D3D11_UNORDERED_ACCESS_VIEW_DESC uavDesc = {
            .Format = desc.Format,
            .ViewDimension = D3D11_UAV_DIMENSION_TEXTURE2D,
            .Texture2D = {.MipSlice = 0}};
        hr = session->device->lpVtbl->CreateUnorderedAccessView(session->device, (ID3D11Resource *)shade, &uavDesc, &uav);
        if (FAILED(hr))
        {
            goto release_frame;
        }
        session->context->lpVtbl->CSSetShader(session->context, session->shader, NULL, 0);
        session->context->lpVtbl->CSSetShaderResources(session->context, 0, 1, &view);
        session->context->lpVtbl->CSSetUnorderedAccessViews(session->context, 0, 1, &uav, NULL);
        UINT groupX = (*width + 7) / 8;
        UINT groupY = (*height + 7) / 8;
        session->context->lpVtbl->Dispatch(session->context, groupX, groupY, 1);

        desc.MipLevels = 1;
        desc.ArraySize = 1;
        desc.Format = DXGI_FORMAT_B8G8R8A8_UNORM;
        desc.SampleDesc.Count = 1;
        desc.SampleDesc.Quality = 0;
        desc.Usage = D3D11_USAGE_STAGING;
        desc.BindFlags = 0;
        desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
        desc.MiscFlags = 0;
    }
    else
    {
        desc.Usage = D3D11_USAGE_STAGING;
        desc.BindFlags = 0;
        desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
        desc.MiscFlags = 0;
    }

    hr = session->device->lpVtbl->CreateTexture2D(session->device, &desc, NULL, &staging);
    if (FAILED(hr))
    {
        goto release_frame;
    }

    if (session->rgb != NULL && session->rgb)
    {
        session->context->lpVtbl->CopyResource(session->context, (ID3D11Resource *)staging, (ID3D11Resource *)shade);
    }
    else
    {
        session->context->lpVtbl->CopyResource(session->context, (ID3D11Resource *)staging, (ID3D11Resource *)texture);
    }

    hr = session->context->lpVtbl->Map(session->context, (ID3D11Resource *)staging, 0, D3D11_MAP_READ, 0, &mapped);
    if (FAILED(hr))
    {
        goto release_frame;
    }

    *bits = (BYTE *)malloc(desc.Width * desc.Height * 4);
    for (int row = 0; row < *height; ++row)
    {
        memcpy(
            *bits + row * (int)(*width) * 4,
            (BYTE *)mapped.pData + ((y + row) * mapped.RowPitch) + x * 4,
            (int)(*width) * 4);
    }

release_frame:
    if (session->duplication != NULL && session->duplication->lpVtbl != NULL)
    {
        session->duplication->lpVtbl->ReleaseFrame(session->duplication);
    }
    if (resource != NULL && resource->lpVtbl != NULL)
    {
        resource->lpVtbl->Release(resource);
        resource = NULL;
    }
    if (staging != NULL && staging->lpVtbl != NULL)
    {
        staging->lpVtbl->Release(staging);
        staging = NULL;
    }
    if (shade != NULL && shade->lpVtbl != NULL)
    {
        shade->lpVtbl->Release(shade);
        shade = NULL;
    }
    if (texture != NULL && texture->lpVtbl != NULL)
    {
        texture->lpVtbl->Release(texture);
        texture = NULL;
    }
    if (view != NULL && view->lpVtbl != NULL)
    {
        view->lpVtbl->Release(view);
        view = NULL;
    }
    if (uav != NULL && uav->lpVtbl != NULL)
    {
        uav->lpVtbl->Release(uav);
        uav = NULL;
    }

    return hr;
}

HRESULT CompileShader(DirectX11Session *session)
{
    HRESULT hr = S_OK;
    ID3DBlob *blob = NULL;
    ID3DBlob *error = NULL;

    hr = D3DCompile(
        shader_code,
        strlen(shader_code),
        NULL,
        NULL,
        NULL,
        "CSMain",
        "cs_5_0",
        D3DCOMPILE_DEBUG | D3DCOMPILE_SKIP_OPTIMIZATION,
        0,
        &blob,
        &error);
    if (FAILED(hr))
    {
        goto release;
    }

    hr = session->device->lpVtbl->CreateComputeShader(session->device, blob->lpVtbl->GetBufferPointer(blob), blob->lpVtbl->GetBufferSize(blob), NULL, &session->shader);

release:
    if (blob != NULL && blob->lpVtbl != NULL)
    {
        blob->lpVtbl->Release(blob);
        blob = NULL;
    }
    if (error != NULL && error->lpVtbl != NULL)
    {
        error->lpVtbl->Release(error);
        error = NULL;
    }

    return hr;
}

HRESULT CountAdapters(UINT *count)
{
    HRESULT hr = S_OK;

    IDXGIFactory1 *factory = NULL;
    IDXGIAdapter *adapter = NULL;

    hr = CreateDXGIFactory1(&IID_IDXGIFactory, (void **)&factory);
    if (FAILED(hr))
    {
        goto release;
    }

    *count = 0;
    for (UINT i = 0;; i++)
    {
        hr = factory->lpVtbl->EnumAdapters(factory, i, &adapter);
        if (hr == DXGI_ERROR_NOT_FOUND)
        {
            break;
        }
        if (FAILED(hr))
        {
            goto release;
        }
        *count = *count + 1;
    }

release:
    if (factory != NULL && factory->lpVtbl != NULL)
    {
        factory->lpVtbl->Release(factory);
        factory = NULL;
    }
    if (adapter != NULL && adapter->lpVtbl != NULL)
    {
        adapter->lpVtbl->Release(adapter);
        adapter = NULL;
    }

    return hr;
}

HRESULT EnumOutputs(UINT adapter_index, UINT index, INT *left, INT *right, INT *top, INT *bottom, WCHAR *device_name, UINT name_size)
{
    HRESULT hr = S_OK;

    IDXGIFactory1 *factory = NULL;
    IDXGIAdapter *adapter = NULL;
    IDXGIOutput *output = NULL;

    DXGI_OUTPUT_DESC desc;

    hr = CreateDXGIFactory1(&IID_IDXGIFactory, (void **)&factory);
    if (FAILED(hr))
    {
        goto release;
    }
    hr = factory->lpVtbl->EnumAdapters(factory, adapter_index, &adapter);
    if (FAILED(hr))
    {
        goto release;
    }
    hr = adapter->lpVtbl->EnumOutputs(adapter, index, &output);
    if (FAILED(hr))
    {
        goto release;
    }
    output->lpVtbl->GetDesc(output, &desc);
    if (!desc.AttachedToDesktop)
    {
        hr = E_UNEXPECTED;
        goto release;
    }

    *left = desc.DesktopCoordinates.left;
    *right = desc.DesktopCoordinates.right;
    *top = desc.DesktopCoordinates.top;
    *bottom = desc.DesktopCoordinates.bottom;
    wcsncpy_s(device_name, name_size, desc.DeviceName, _TRUNCATE);

release:
    if (factory != NULL && factory->lpVtbl != NULL)
    {
        factory->lpVtbl->Release(factory);
        factory = NULL;
    }
    if (adapter != NULL && adapter->lpVtbl != NULL)
    {
        adapter->lpVtbl->Release(adapter);
        adapter = NULL;
    }
    if (output != NULL && output->lpVtbl != NULL)
    {
        output->lpVtbl->Release(output);
        output = NULL;
    }

    return hr;
}

// int main()
// {
//     HRESULT hr = S_OK;
//     DirectX11Session session;
//     UINT display = 0;
//     printf("start\n");
//     hr = NewDirectX11Session(&session, display);
//     if (SUCCEEDED(hr))
//     {
//         printf("success\n");
//     }
//     else
//     {
//         printf("failed\n");
//     }

//     CloseDirectX11Session(&session);
// }